import rclpy
from rclpy.node import Node
from base_interfaces_demo.msg import SensorData
from pymongo import MongoClient  # 新增 MongoDB 驱动
import random
import datetime

class PubDataNode(Node):
    def __init__(self, name):
        super().__init__(name)
        
        # ----------------- 数据发布部分 -----------------
        self.sensor_publisher = self.create_publisher(SensorData, "/sensor_data", 10)
        self.timer = self.create_timer(1.0, self.timer_callback)
        
        # ----------------- 数据库连接部分 -----------------
        self.client = MongoClient('192.168.85.121', 27017)  # 替换为你的 MongoDB IP
        self.db = self.client['sensor_db']
        self.collection = self.db['sensor_data']
        self.get_logger().info("成功连接 MongoDB 数据库！")
        
        self.get_logger().info("传感器节点已启动 (发布+存储)")

    def timer_callback(self):
        # 生成模拟传感器数据
        msg = SensorData()
        msg.temperature = random.uniform(-100.0, 120.0)
        msg.humidity = random.uniform(30.0, 80.0)
        
        # 发布到 ROS 话题
        self.sensor_publisher.publish(msg)
        self.get_logger().info(f"发布数据: 温度 {msg.temperature:.1f}°C, 湿度 {msg.humidity:.1f}%")
        
        # 存储到 MongoDB（新增部分）
        try:
            data_doc = {
                'temperature': msg.temperature,
                'humidity': msg.humidity,
                'timestamp': datetime.datetime.now()
            }
            self.collection.insert_one(data_doc)
            self.get_logger().info("数据已存入数据库", throttle_duration_sec=5)  # 限流日志
        except Exception as e:
            self.get_logger().error(f"数据库写入失败: {str(e)}")

    def __del__(self):
        """节点销毁时自动关闭数据库连接"""
        if hasattr(self, 'client'):
            self.client.close()
            self.get_logger().info("MongoDB 连接已关闭")

def main(args=None):    
    rclpy.init(args=args)    
    node = PubDataNode("pub_data_node")  
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == "__main__":
    main()